import isa
from scanner import *

def check_token_type(token : Token, tartget, message):
    if token.token_type != tartget:
        raise Exception(message)

def assemble_register(token : Token, bias : int, message : str) -> int:
    if (token.token_type != TokenType.REGISTER):
        raise Exception(message)
    else:
        rd_index = int(token.context[1 : ])
        return (rd_index & 0b11111) << bias

def assemble_rd(token : Token) -> int:
    return assemble_register(token, 7, '[error] expect a resgiter as rd')
    
def assemble_rs1(token : Token) -> int:
    return assemble_register(token, 15, '[error] expect a resgiter as rd')

def assemble_rs2(token : Token) -> int:
    return assemble_register(token, 20, '[error] expect a resgiter as rd')

def assemble_itype_imm(token : Token) -> int:
    if (token.token_type != TokenType.NUMBER):
        raise Exception('[error] expect a number as i-type imm')
    else:
        imm = int(token.context)
        return (imm & 0xFFF) << 20
        
def assemble_stype_or_btype_imm(token : Token) -> int:
    if (token.token_type != TokenType.NUMBER):
        raise Exception('[error] expect a number as r-type/s-type imm')
    else:
        imm = int(token.context)
        value = 0
        value += (imm & 0b11111) << 7
        value += ((imm & 0xFFF) >> 5) << 20
        return value

def assemble_jtype_or_utype_imm(token : Token) -> int:
    if (token.token_type != TokenType.NUMBER):
        raise Exception('[error] expect a number as j-type/u-type imm')
    else:
        imm = int(token.context)
        return (imm & 0xFFFFF) << 20

class InstrAssembler:
    def assemble(self, instr : list[Token]) -> int:
        pass

class ITypeAssembler(InstrAssembler):
    # addi x1, x0, 2
    def assemble(self, instr : list[Token]) -> int:
        # opcode
        value : int = isa.ITYPE_OPCODE
        command = instr[0].context

        # alu opcode
        alu_opcode = isa.alu_opcode_dict[command]
        value += (alu_opcode & 0b111) << 12
        
        # rd: instr[1] must be a register token
        value += assemble_rd(instr[1])
        check_token_type(instr[2], TokenType.COMMA, '[error] expect \',\'')

        # rs1: instr[3] must be a register token
        value += assemble_rs1(instr[3])
        check_token_type(instr[4], TokenType.COMMA, '[error] expect \',\'')

        # imm
        value += assemble_itype_imm(instr[5])

        return value & 0xFFFFFFFF

class RTypeAssembler(InstrAssembler):
    # add x1, x2, x3
    def assemble(self, instr : list[Token]) -> int:
        # opcode
        value : int = isa.RTYPE_OPCODE
        command = instr[0].context

        # alu opcode
        value += (isa.alu_opcode_dict[command] << 12) & 0b111
        
        # rd: instr[1] must be a register token
        value += assemble_rd(instr[1])
        check_token_type(instr[2], TokenType.COMMA, '[error] expect \',\'')

        # rs1: instr[3] must be a register token
        value += assemble_rs1(instr[3])
        check_token_type(instr[4], TokenType.COMMA, '[error] expect \',\'')

        # rs2: instr[5] must be a register token
        value += assemble_rs2(instr[5])

        return value & 0xFFFFFFFF

class STypeAssembler(InstrAssembler):
    # sw rs2, 12(rs1)
    def assemble(self, instr : list[Token]) -> int:
        # opcode
        value : int = isa.STYPE_OPCODE

        # rs2
        value += assemble_rs2(instr[1])
        check_token_type(instr[2], TokenType.COMMA, '[error] expect \',\'')

        # imm
        value += assemble_stype_or_btype_imm(instr[3])
        check_token_type(instr[4], TokenType.LEFT_PARE, '[error] expect \'(\'')

        # rs1
        value += assemble_rs1(instr[5])
        check_token_type(instr[6], TokenType.RIGHT_PARE, '[error] expect \'(\'')

        return value & 0xFFFFFFFF


class BTypeAssembler(InstrAssembler):
    # beq rs1, rs2, 12
    def assemble(self, instr : list[Token]) -> int:
        # opcode
        value : int = isa.BTYPE_OPCODE
        command = instr[0].context

        # branch type
        value += (isa.cmp_opcode_dict[command] << 12) & 0b000

        # rs1
        value += assemble_rs2(instr[1])
        check_token_type(instr[2], TokenType.COMMA, '[error] expect \',\'')

        # rs2
        value += assemble_rs1(instr[3])
        check_token_type(instr[4], TokenType.COMMA, '[error] expect \'(\'')

        # imm
        value += assemble_stype_or_btype_imm(instr[5])

        return value & 0xFFFFFFFF

class JTypeAssembler(InstrAssembler):
    # jal r1, 12
    def assemble(self, instr : list[Token]) -> int:
        value : int = isa.JTYPE_OPCODE

        # rd
        value += assemble_rd(instr[1])
        check_token_type(instr[2], TokenType.COMMA, '[error] expect \',\'')

        # imm
        value += assemble_jtype_or_utype_imm(instr[3])

        return value & 0xFFFFFFFF
        
class UTypeAssembler(InstrAssembler):
    # lui r1, 12
    def assemble(self, instr : list[Token]) -> int:
        value : int = isa.UTYPE_OPCODE

        # rd
        value += assemble_rd(instr[1])
        check_token_type(instr[2], TokenType.COMMA, '[error] expect \',\'')

        # imm
        value += assemble_jtype_or_utype_imm(instr[3])

        return value & 0xFFFFFFFF

class LwAssembler(InstrAssembler):
    # lw x1, 12(x4)
    def assemble(self, instr : list[Token]) -> int:
        # opcode
        value : int = isa.LW_OPCODE
        
        # rd: instr[1] must be a register token
        value += assemble_rd(instr[1])
        check_token_type(instr[2], TokenType.COMMA, '[error] expect \',\'')

        # imm
        value += assemble_itype_imm(instr[3])
        check_token_type(instr[4], TokenType.LEFT_PARE, '[error] expect \'(\'')

        # rs1
        value += assemble_rs1(instr[5])
        check_token_type(instr[6], TokenType.RIGHT_PARE, '[error] expect \')\'')

        return value & 0xFFFFFFFF

assemblers : dict[TokenType, InstrAssembler] = {
    TokenType.ITYPE_COMMAND : ITypeAssembler(),
    TokenType.RTYPE_COMMAND : RTypeAssembler(),
    TokenType.STYPE_COMMAND : STypeAssembler(),
    TokenType.BTYPE_COMMAND : BTypeAssembler(),
    TokenType.JTYPE_COMMAND : JTypeAssembler(),
    TokenType.UTYPE_COMMAND : UTypeAssembler(),
    TokenType.LW_COMMAND : LwAssembler(),
}
